Pretraining Transformers with Quantized Softmax in Attention
预训练中量化注意力 Softmax：校准与直通梯度设计的影响
🧭 低精度注意力 Softmax 训练 ★ 深度解读
University of Illinois Urbana-Champaign
系统研究量化 softmax 前向与反向规则对预训练稳定性的影响，给出可复现的设计选择
📌 工作参考价值（4 点）
- 工业推荐/Agent 里若对 Transformer 做 QAT 或量化 softmax 部署，不要用 .detach() 隔离 calibration 统计量（row max/min）；必须让 row extrema 反向传播，否则训练前 25-30M tokens 看似正常，之后 loss 突然恶化，末尾可差 0.65-3.07 nats。零和投影不能修复。
- 做 hard rounding 量化时，优先 Prob-STE（在归一化后概率上放 STE），或校准从 MinMax 换成 fixed window to max（FWM）。在 K=4、124M/2.5B tokens 下 MinMax+Weight-STE 差 softmax 0.89 nats，而 FWM+Prob-STE 只差 0.019；K=16 时 FWM+Weight-STE 差 0.004。
- 如果允许可微实现，优先 LERP（分段线性插值）而非 Nearest hard rounding；K=4/16/32 在 2.5B tokens 都能收敛到 softmax 0.005 nats 以内。
- 评估时必须用训练时的 native forward；训练后替换成 exact softmax 会使 gap 扩大，LERP K=2 从 +0.005 跳到 +1.814 nats，说明权重已经依赖训练算子。